/**
 * Copyright (c) iwindplus Technologies Co., Ltd.2024-2030, All rights reserved.
 */

package com.iwindplus.boot.web.filter;

import cn.hutool.core.map.MapUtil;
import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.http.HtmlUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.iwindplus.boot.domain.constant.CommonConstant.ExceptionConstant;
import com.iwindplus.boot.util.JacksonUtil;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

/**
 * 参数去除空格，实现对request body的二次读取，防止夸脚本工具攻击.
 *
 * @author zengdegui
 * @since 2020/4/3
 */
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    @Override
    public String getQueryString() {
        String value = super.getQueryString();
        if (value == null) {
            return null;
        }
        return CharSequenceUtil.trim(HtmlUtil.filter(value));
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        if (value == null) {
            return null;
        }
        return CharSequenceUtil.trim(HtmlUtil.filter(value));
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values == null) {
            return null;
        }
        return Stream.of(values)
            .filter(ObjectUtil::isNotEmpty)
            .map(value -> CharSequenceUtil.trim(HtmlUtil.filter(value)))
            .toArray(String[]::new);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameters = super.getParameterMap();
        if (MapUtil.isEmpty(parameters)) {
            return parameters;
        }
        return parameters.entrySet().stream()
            .filter(entry -> ObjectUtil.isNotEmpty(entry.getKey()))
            .collect(Collectors.toMap(
                Map.Entry::getKey,
                entry -> Stream.of(entry.getValue())
                    .filter(ObjectUtil::isNotEmpty)
                    .map(value -> CharSequenceUtil.trim(HtmlUtil.filter(value)))
                    .toArray(String[]::new)
            ));
    }

    @Override
    public Enumeration<String> getParameterNames() {
        Enumeration<String> enumeration = super.getParameterNames();
        return this.getEnumeration(enumeration);
    }

    @Override
    public Object getAttribute(String name) {
        Object value = super.getAttribute(name);
        return this.getValue(value);
    }

    @Override
    public Enumeration<String> getAttributeNames() {
        Enumeration<String> enumeration = super.getAttributeNames();
        return this.getEnumeration(enumeration);
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return value == null ? null : CharSequenceUtil.trim(HtmlUtil.filter(value));
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        Enumeration<String> headers = super.getHeaders(name);
        return this.getEnumeration(headers);
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        ServletInputStream inputStream = super.getInputStream();
        if (!this.isJsonRequest()) {
            return inputStream;
        }
        String json = this.getJson(inputStream);
        if (ObjectUtil.isEmpty(json)) {
            return inputStream;
        }
        byte[] jsonBytes = this.parseJsonToBytes(json);
        return this.getServletInputStream(jsonBytes);
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream(), Charset.defaultCharset()));
    }

    private Enumeration<String> getEnumeration(Enumeration<String> headers) {
        List<String> list = new ArrayList<>(10);
        while (headers.hasMoreElements()) {
            String element = headers.nextElement();
            if (ObjectUtil.isNotEmpty(element)) {
                list.add(CharSequenceUtil.trim(HtmlUtil.filter(element)));
            }
        }
        return Collections.enumeration(list);
    }

    private ServletInputStream getServletInputStream(byte[] jsonBytes) {
        final ByteArrayInputStream bis = new ByteArrayInputStream(jsonBytes);
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return true;
            }

            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public int available() {
                return jsonBytes.length;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
                // Do nothing
            }

            @Override
            public int read() {
                return bis.read();
            }
        };
    }

    private boolean isJsonRequest() {
        String header = this.getHeader(HttpHeaders.CONTENT_TYPE);
        return CharSequenceUtil.startWithIgnoreCase(header, MediaType.APPLICATION_JSON_VALUE);
    }

    private String getJson(ServletInputStream inputStream) {
        StringBuilder stringBuilder = new StringBuilder();
        try (InputStreamReader inputStreamReader = new InputStreamReader(inputStream, Charset.defaultCharset());
            BufferedReader bufferedReader = new BufferedReader(inputStreamReader)) {
            String line;
            while ((line = bufferedReader.readLine()) != null) {
                stringBuilder.append(CharSequenceUtil.trim(line));
            }
        } catch (IOException ex) {
            log.error(ExceptionConstant.IO_EXCEPTION, ex);
        }
        return stringBuilder.toString();
    }

    private byte[] parseJsonToBytes(String json) {
        final JsonNode jsonNode = JacksonUtil.parseTree(json);
        if (jsonNode.isArray()) {
            List<Map<String, Object>> result = new ArrayList<>(10);
            JacksonUtil.parseList(json).forEach(data -> result.add(this.getFilterMap(data)));
            return JacksonUtil.toJsonBytes(result);
        } else if (jsonNode.isObject()) {
            Map<String, Object> parameters = JacksonUtil.parseObject(json, new TypeReference<>() {
            });
            return JacksonUtil.toJsonBytes(this.getFilterMap(parameters));
        }

        return new byte[0];
    }

    private Map<String, Object> getFilterMap(Map<String, Object> parameters) {
        return parameters.entrySet().stream()
            .filter(entry -> ObjectUtil.isNotEmpty(entry.getKey()) && entry.getValue() != null)
            .collect(Collectors.toMap(
                Map.Entry::getKey,
                entry -> this.getValue(entry.getValue())
            ));
    }

    private Object getValue(Object value) {
        return value instanceof String str ? CharSequenceUtil.trim(HtmlUtil.filter(str)) : value;
    }

}
